837B - Flag of Berland - CodeForces Solution

#include <bits/stdc++.h>

using namespace std;
using ll = long long;

bool ok(vector<string> v){
    int n = v.size();
    int m = v[0].size();

    for(int i = 0; i < n; i++){
        for(int j = 1; j < m; j++){
            if(v[i][j] != v[i][j-1]) return 0;
        }
    }

    vector<char> v1;
    vector<int> w;
    for(int i = 0; i < n; i++){
        if(i == 0 || v[i][0] != v[i-1][0]){
            v1.push_back(v[i][0]);
            w.push_back(1);
        } else w.back()++;
    }
    if(v1.size() != 3) return 0;
    if(w[0] != w[1] || w[1] != w[2]) return 0;

    set<char> s(v1.begin(), v1.end());
    if(s.size() != 3) return 0;

    return 1;
}

vector<string> trans(vector<string> v){
    vector<string> res;
    int n = v.size(), m = v[0].size();
    for(int j = 0; j < m; j++){
        string cur;
        for(int i = 0; i < n; i++){
            cur += v[i][j];
        }
        res.push_back(cur);
    }
    return res;
}

int main() {

    int n, m;
    cin >> n >> m;

    vector<string> v(n);
    for(string &x : v) cin >> x;

    if(ok(v) || ok(trans(v))) cout << "YES";
    else cout << "NO";

    return 0;
}